/**
 * 
 */
package edu.umd.clip.smoothing;

import java.io.PrintStream;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class KNModSmoother extends Smoother {
	private double D1;
	private double D2;
	private double D3;
	private final double backoffProb;
	private final double revTotalCount;
	
	public KNModSmoother(SmootherProducer producer) {
		this(producer, null);
	}
	
	public KNModSmoother(SmootherProducer producer, PrintStream verboseOutput) {
		super(producer);
		revTotalCount = 1.0 / totalCount;
		int counts[] = producer.getCounts();
		int N1 = counts[1];
		int N2 = counts[2];
		int N3 = counts[3];
		int N4 = counts[4];
		
		if (N1 == 0 || N2 == 0 || N3 == 0 || N4 == 0) {
			D1 = D2 = D3 = 0;
			backoffProb = Double.NaN;
		} else {
			double D = (double) N1 / (N1 + 2*N2);
			D1 = 1 - 2 * D * N2 / N1;
			D2 = 2 - 3 * D * N3 / N2;
			D3 = 3 - 4 * D * N4 / N3;
			
			if (D1 > 0 && D2 >= D1 && D3 >= D2) {
				int c3 = 0;
				for(int i = 3; i<counts.length; ++i) {
					c3 += counts[i];
				}
				double totalDiscount = D1 * counts[1] + D2 * counts[2] + D3 * c3;
				backoffProb = totalDiscount / producer.getTotalCount();
			} else {
				//D1 = D2 = D3 = 0;
				backoffProb = Double.NaN;				
			}
		}
		
		if (verboseOutput != null) {
			verboseOutput.printf("KNMod: N1=%d, N2=%d, N3=%d, N4=%d, D1=%g, D2=%g, D3=%g, backoff=%g, count=%d\n", 
					N1, N2, N3, N4, D1, D2, D3, backoffProb, producer.getTotalCount());
		}
		
	}
	
	public KNModSmoother(SmootherProducer producer, double D1, double D2, double D3) {
		super(producer);
		
		this.D1 = D1;
		this.D2 = D2;
		this.D3 = D3;
		
		revTotalCount = 1.0 / totalCount;
		int counts[] = producer.getCounts();
		int c3 = 0;
		for(int i = 3; i<counts.length; ++i) {
			c3 += counts[i];
		}
		double totalDiscount = D1 * counts[1] + D2 * counts[2] + D3 * c3;
		backoffProb = totalDiscount / producer.getTotalCount();
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.smoothing.Smoother#getBackoffProb()
	 */
	public double getBackoffProb() {
		return backoffProb;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.smoothing.Smoother#getProb(long)
	 */
	public double getProb(long count) {
		if (count < minCount) return 0.0; 
		if (count == 1) {
			return (count - D1) * revTotalCount;
		} else if (count == 2) {
			return (count - D2) * revTotalCount;
		} else if (count > 2) {
			return (count - D3) * revTotalCount;
		}
		return 0;
	}
	public double getD1() {
		return D1;
	}
	public double getD2() {
		return D2;
	}
	public double getD3() {
		return D3;
	}

	@Override
	public String toString() {
		return "KNModSmoother [D1=" + D1 + ", D2=" + D2 + ", D3=" + D3
				+ ", backoffProb=" + backoffProb + ", totalCount=" + totalCount
				+ "]";
	}
}
